import asyncio
import time
from arq.connections import RedisSettings
import os
import uuid
from datetime import datetime
from config import redis_url, arq_worker_name, arq_queue_name, redis_host, redis_port, redis_db, redis_conn_timeout, redis_conn_retry

# 全局状态存储
task_status = {}
task_paused = {}

async def startup(ctx):
    """Worker启动时执行"""
    ctx["startup_time"] = time.time()
    print(f"{arq_worker_name} 已启动")

async def shutdown(ctx):
    """Worker关闭时执行"""
    print(f"{arq_worker_name} 运行了 {time.time() - ctx['startup_time']:.2f} 秒")

async def process_file(ctx, file_path: str, file_id: str, task_id: str, original_filename: str, file_size: int):
    """
    ARQ任务处理器 - 处理文件上传
    """
    try:
        # 确保task_id和file_id存在于task_status中
        if task_id not in task_status or file_id not in task_status[task_id]["files"]:
            print(f"任务ID {task_id} 或文件ID {file_id} 不存在于任务状态中")
            return {"success": False, "file_id": file_id, "error": "任务状态丢失"}
        
        # 更新任务开始状态
        task_status[task_id]["files"][file_id].update({
            "status": "processing",
            "progress": 0
        })
        
        # 生成文件URL（确保URL正确设置）
        file_url = task_status[task_id]["files"][file_id].get("file_url", "")
        if not file_url:
            file_url = f"http://localhost:8005/static/uploads/{os.path.basename(file_path)}"
            task_status[task_id]["files"][file_id]["file_url"] = file_url
        
        print(f"开始处理文件: {file_id}, 任务: {task_id}")
        
        # 减少处理步骤和等待时间
        total_steps = 5
        for step in range(1, total_steps + 1):
            # 检查任务是否已被暂停或删除
            if task_id in task_paused and task_paused[task_id]:
                task_status[task_id]["files"][file_id].update({
                    "status": "paused",
                    "paused_at": datetime.now().isoformat(),
                    "paused_progress": task_status[task_id]["files"][file_id]["progress"]
                })
                update_total_progress(task_id)
                return {"success": False, "file_id": file_id, "status": "paused"}
            
            # 检查任务是否被删除
            if task_id not in task_status:
                return {"success": False, "file_id": file_id, "status": "deleted"}
                
            # 更新当前文件进度 - 快速增加进度
            progress = int((step / total_steps) * 100)
            task_status[task_id]["files"][file_id]["progress"] = progress
            print(f"文件 {file_id} 进度更新: {progress}%")
            
            # 更新总体进度
            update_total_progress(task_id)
            
            # 缩短等待时间
            await asyncio.sleep(0.1)
        
        # 处理完成，更新状态
        task_status[task_id]["files"][file_id].update({
            "status": "completed",
            "progress": 100,
            "completed_at": datetime.now().isoformat()
        })
        
        # 更新总体进度
        update_total_progress(task_id)
        
        # 检查所有文件是否都已完成
        all_completed = all(file_info["status"] == "completed" 
                           for file_info in task_status[task_id]["files"].values())
        
        if all_completed:
            task_status[task_id]["status"] = "completed"
            task_status[task_id]["completed_at"] = datetime.now().isoformat()
        
        print(f"文件 {file_id} 处理完成")
        return {
            "success": True, 
            "file_id": file_id, 
            "status": "completed",
            "file_url": file_url
        }
    
    except Exception as e:
        # 处理失败，更新状态
        if task_id in task_status and file_id in task_status[task_id]["files"]:
            task_status[task_id]["files"][file_id].update({
                "status": "failed",
                "error": str(e)
            })
            update_total_progress(task_id)
        print(f"文件 {file_id} 处理失败: {str(e)}")
        return {"success": False, "file_id": file_id, "error": str(e)}

def update_total_progress(task_id: str):
    """更新任务总体进度"""
    if task_id not in task_status:
        return
    
    files_info = task_status[task_id]["files"]
    if not files_info:
        return
    
    # 计算总体进度为所有文件进度的平均值
    total_progress = sum(file_info["progress"] for file_info in files_info.values()) / len(files_info)
    task_status[task_id]["progress"] = int(total_progress)
    print(f"任务 {task_id} 总进度更新: {task_status[task_id]['progress']}%")

class WorkerSettings:
    """ARQ Worker 配置"""
    # ARQ中RedisSettings的正确参数设置
    redis_settings = RedisSettings(
        host=redis_host,
        port=redis_port,
        database=redis_db,
        conn_timeout=redis_conn_timeout,
        conn_retries=redis_conn_retry,
        conn_retry_delay=1.0
    )
    functions = [process_file]
    on_startup = startup
    on_shutdown = shutdown
    queue_name = arq_queue_name
    worker_name = arq_worker_name
    max_jobs = 10
    job_timeout = 300 